#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@IDE     ：PyCharm 
@Author  ：kiway
@Date    ：2024/10/12 20:15 
"""

import lmdb
import os
import pandas as pd
import datetime
import csv

if not os.path.exists("data"):
    os.makedirs("data")


class LmdbHandler:

    def __init__(self, path):
        self.env = lmdb.open(path, map_size=int(0.5 * 1024 * 1024 * 1024))

    def insert(self, sid, name):
        txn = self.env.begin(write=True)
        txn.put(str(sid).encode(), name.encode())
        txn.commit()

    def delete(self, sid):
        txn = self.env.begin(write=True)
        txn.delete(str(sid).encode())
        txn.commit()

    def update(self, sid, name):
        txn = self.env.begin(write=True)
        txn.put(str(sid).encode(), name.encode())
        txn.commit()

    def search(self, sid):
        txn = self.env.begin()
        name = txn.get(str(sid).encode())
        return name

    def display(self):
        txn = self.env.begin()
        cur = txn.cursor()
        for key, value in cur:
            print(key, value)

    def search_all(self):
        data = {}
        txn = self.env.begin()
        cur = txn.cursor()

        for key, value in cur:
            data[key] = eval(value)
        return data

    def clear(self):
        txn = self.env.begin()
        cur = txn.cursor()
        for key, value in cur:
            txn.delete(str(key).encode())

        txn.commit()


class TradeRecord(LmdbHandler):

    def __init__(self):
        super().__init__("data/traders")
        self.clear()


class BalanceRecord(LmdbHandler):

    def __init__(self):
        super().__init__("data/balance")


class PositionRecord(LmdbHandler):
    def __init__(self):
        super().__init__("data/position")


class OrderRecord(LmdbHandler):
    def __init__(self):
        super().__init__("data/order")


class LatestRecord(LmdbHandler):
    def __init__(self):
        super().__init__("data/sharedMem")

    def clear(self):
        txn = self.env.begin(write=True)
        txn.put(str("latest_pos").encode(), str([]).encode())
        txn.put(str("latest_trade").encode(), str([]).encode())
        txn.put(str("latest_balance").encode(), str({"info": []}).encode())
        txn.put(str("time_cut_forward").encode(), str("0000-00-00 00:00:00").encode())

        t = pd.to_datetime(datetime.datetime.now()).floor('min')
        txn.put(str("time_cut_backward").encode(), str(t).encode())

        txn.commit()

    def position_info(self):
        return eval(self.search("latest_pos"))

    def trade_info(self):
        return eval(self.search("latest_trade"))

    def balance_info(self):
        return eval(self.search("latest_balance"))

    def time_cut_forward(self):
        return self.search("time_cut_forward").decode()

    def time_cut_backward(self):
        return self.search("time_cut_backward").decode()


class MarketSave:

    def __init__(self):

        self.parent_path = r"data/market_data"
        self.child_bar_path = r"data/market_data/bar"
        self.child_aggTrade_path = r"data/market_data/aggTrade"
        self.columns = None
        self.check_path()
        self.get_columns()
        self.aggTrade_data_buffer = []

    def check_path(self):
        if not os.path.exists("data"):
            os.makedirs("data")

        if not os.path.exists(self.parent_path):
            os.makedirs(self.parent_path)

        if not os.path.exists(self.child_bar_path):
            os.makedirs(self.child_bar_path)

        if not os.path.exists(self.child_aggTrade_path):
            os.makedirs(self.child_aggTrade_path)

    def get_columns(self):

        self.columns = {'spot_aggTrade': ["e", "E", "s", "a", "p", "q", 'f', "l", "T", "m", "M"],
                        'future_aggTrade': ["e", "E", "s", "a", "p", "q", "f", "l", "T", "m"],
                        'spot_bar': ["e", "E", "s", "k", "k_t", "k_T", "k_s", "k_i", "k_f", "k_L", "k_0",
                                     "k_c", "k_h", "k_l", "k_v", "k_n", "k_x", "k_q", "k_V", "k_Q", "k_B"],

                        'future_bar': ["e", "E", "s", "k_t", "k_T", "k_s", "k_i", "k_f", "k_L", "k_o",
                                       "k_c", "k_h", "k_l", "k_v", "k_n", "k_x", "k_q", "k_v", "k_Q", "k_B"]}

    def extract_values(self, data):
        values = []
        for value in data.values():
            if isinstance(value, dict):  # 如果值是字典，递归提取其值
                values.extend(self.extract_values(value))  # 递归调用
            else:
                values.append(value)  # 如果不是字典，直接添加到列表中
        return values

    def ensure_csv_file(self, file_name, type_):
        """

        :param file_name:
        :param type_: spot_aggTrades
        :return:
        """

        if not os.path.exists(file_name):
            with open(file_name, 'w', newline='', encoding='utf-8') as file:
                writer = csv.writer(file)
                # 写入表头
                col = self.columns[type_]
                writer.writerow(col)

    def dist_spot_future(self, symbol):
        if symbol.endswith(":USDT"):  # 期货
            return 1
        else:
            return 0

    def write_to_csv(self, file_name, data, type_):
        if type_ == "spot_aggTrade" or type_ == "future_aggTrade":
            file_path = self.child_aggTrade_path + "/" + file_name + ".csv"

        else:
            file_path = self.child_bar_path + "/" + file_name + '.csv'

        self.ensure_csv_file(file_path, type_)

        with open(file_path, 'a', newline='', encoding='utf-8') as file:
            writer = csv.writer(file)
            writer.writerow(data)

    def save_bar(self, symbol, data):
        tm = data[symbol]["1m"][0][1]
        time_cut = str(pd.to_datetime(tm, unit='ms').floor('min'))[:10]
        dt = data[symbol]["1m"][0]

        if self.dist_spot_future(symbol):
            file_name = time_cut + "_" + symbol.split("/USDT")[0]
            type_ = "future_bar"
            file_name += "_future"
        else:
            file_name = time_cut + "_" + symbol.split("/USDT")[0]
            type_ = "spot_bar"
            file_name += "_spot"

        self.write_to_csv(file_name, dt, type_)

    def save_aggTrade(self, symbol, data):
        tm = data[0]["info"]["E"]
        time_cut = str(pd.to_datetime(tm, unit='ms').floor('min'))[:10]
        dt = list(data[0]["info"].values())
        if self.dist_spot_future(symbol):
            file_name = time_cut + "_" + symbol.split("/USDT")[0]
            type_ = "future_aggTrade"
            file_name += "_future"
        else:
            file_name = time_cut + "_" + symbol.split("/USDT")[0]
            type_ = "spot_aggTrade"
            file_name += "_spot"

        self.write_to_csv(file_name, dt, type_)
